import matplotlib.pyplot as plt
import matplotlib.animation as anime

from hartreefock import *

def do_H2_calculation():
    '''
    计算 H2 电子云变化
    '''
    r = np.linspace(5, 0.1, num=100)
    E = []

    fig = plt.figure(figsize=(10, 10), dpi=144)
    plt.clf()
    # plt.ion()
    def update(t):
        print(t)
        ri = r[t]
        Ei, ei, basis, Cx = HartreeFock(H2(ri), details=True)
        C1 = Cx[:, 0]
        E.append(Ei)

        z = np.linspace(-1.5, ri+1.5, num=1000)
        d1 = eval_cloud(basis, C1, z, 0)
        plt.clf()
        
        ax = plt.gca()
        ax.spines['top'].set_visible(False)
        ax.spines['left'].set_visible(False)
        ax.spines['right'].set_visible(False)
        plt.tick_params(bottom=True, top=False, left=False, right=False)
        plt.xlabel(r'$r/\mathrm{a.u.}$')
        plt.yticks([])
        plt.xlim([-4.5, 4.5])
        plt.ylim([0, 0.8])
        plt.plot(z - ri/2, d1, label='electron cloud')
        # plt.legend()

        plt.axes([0.1, 0.7, 0.25, 0.25])
        ax = plt.gca()
        ax.spines['top'].set_visible(False)
        ax.spines['left'].set_visible(False)
        ax.spines['right'].set_visible(False)
        plt.xlim([0, 5])
        plt.ylim([-1.5, 7.5])
        plt.xlabel(r'$r/\mathrm{a.u.}$')
        plt.yticks([])
        plt.plot(r[0:t+1], E)
        plt.text(2, 3, f'$E_0$ = {Ei:.3f} a.u.')

    # 如果要生成图片，使用这部分代码，把下面那一部分注释掉
    # H2_anime = anime.FuncAnimation(fig, func=update, frames=100, init_func=lambda : plt.clf())
    # H2_anime.save("H2.mp4")
    
    plt.ion()
    for t in range(100):
        update(t)
        plt.pause(0.05)
    plt.ioff()
    plt.show()



if __name__ == '__main__':
    do_H2_calculation()
